#include <iostream>
#include <cmath>
#include <vector>
#include "Armor.h"


double getDistance(const cv::Point2f &c1, const cv::Point2f &c2) {
    return pow(pow(c1.x - c2.x, 2) + pow(c1.y - c2.y, 2), 0.5);
}

int main() {
    cv::Mat frame, img, hsv, mask;
    cv::VideoCapture video;
    video.open("/home/champrin/CLionProjects/evolution-vision-puzzle-3/8.avi");
    namedWindow("video", cv::WINDOW_AUTOSIZE);
    EnemyColor color = BLUE;/*
    while (cv::waitKey(0) == 'i'){
        break;
    }*/
    while (video.read(frame)) {
        cv::Mat channels[3];
        split(frame, channels);
        mask = channels[0] - channels[2];

        threshold(mask, mask, 150, 255, cv::THRESH_BINARY);
        dilate(mask, mask, getStructuringElement(0, cv::Size(3, 3)));

        erode(mask, mask, getStructuringElement(0, cv::Size(3, 3)));
        dilate(mask, mask, getStructuringElement(0, cv::Size(3, 3)));
//        imshow("k", mask);
        std::vector<std::vector<cv::Point>> contours;
        std::vector<cv::RotatedRect> choose;
        findContours(mask, contours, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE, cv::Point());
        double area;
        cv::RotatedRect light;

        for (auto &contour : contours) {
            area = contourArea(contour);
            if (area < 3.6f || area > 40.0f) continue;
//            cout<<area<<endl;
            light = minAreaRect(contour);

            if (
                    ((fabs(light.angle) > 70.0f) && (light.size.width >= light.size.height) && light.size.width > 1.9f)||
                ((fabs(light.angle) < 15.0f) && (light.size.height > light.size.width)) && light.size.height > 7.0f
                ) {
                choose.push_back(light);
            }
        }
//        cout<<contours.size()<<endl;
        cv::RotatedRect tmp1, tmp2;
        for (int i = 0; i < choose.size(); ++i) {
            tmp1 = choose.at(i);
            for (int j = i + 1; j < choose.size(); ++j) {
                tmp2 = choose.at(j);

                double distance = getDistance(tmp1.center, tmp2.center);

                if (distance <= 40.0 || distance>= 165.0)
                    continue;
                if (std::max(tmp1.size.height, tmp1.size.width) / std::max(tmp2.size.height, tmp2.size.width) > 3.8 ||
                        std::max(tmp1.size.height, tmp1.size.width) / std::max(tmp2.size.height, tmp2.size.width) < 0.26)
                    continue;

                if (abs(tmp1.center.y - tmp2.center.y) > 5)
                    continue;
                if (abs(tmp1.angle - tmp2.angle) > 6)
                    continue;

                Armor a1;
                a1.SetEnemyColor(color);
                a1.Transform2P(tmp1, tmp2);
                a1.DrawRectangle(frame, cv::Scalar(0, 255, 0));
                a1.SolvePNP();
                float dis = a1.GetDistance();
                std::cout<<dis<<std::endl;
                cv::putText(frame, std::to_string(dis/1000)+"m", a1.p1, cv::FONT_HERSHEY_SIMPLEX, 1.0f, cv::Scalar(0, 255, 0));
                break;
            }
        }
        imshow("video", frame);
        cv::waitKey(1000 / video.get(cv::CAP_PROP_FPS));
    }
    cv::waitKey(0);
    return 0;
}